package com.xiaoyu.awakening.ai.index.config;

import cn.hutool.core.thread.NamedThreadFactory;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.Resource;
import lombok.extern.log4j.Log4j2;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.concurrent.*;

/**
 * 索引线程池配置
 *
 * @author wangxiaoyu
 * @date 2025/6/25 16:03
 * @since v1.0
 */
@Configuration
@EnableConfigurationProperties(AwkThreadPoolProperties.class)
@Log4j2
public class AwkThreadPoolConfig {
    /**
     * 定义线程池名称
     */
    public static final String INDEX_TASK_THREAD_POOL_NAME = "indexTaskThreadPool";
    public static final String OTHER_TASK_THREAD_POOL_NAME = "otherTaskThreadPool";
    public static final ConcurrentHashMap<String, ExecutorService> AWK_THREAD_POOL = new ConcurrentHashMap<>();
    private static final ScheduledThreadPoolExecutor MONITOR = new ScheduledThreadPoolExecutor(1, new NamedThreadFactory("awk-thread-monitor", true));


    /**
     * 注入用户自定义线程池配置
     */
    @Resource
    private AwkThreadPoolProperties threadPoolProperties;

    /**
     * 定义一个线程池
     */
    @Bean(name = INDEX_TASK_THREAD_POOL_NAME)
    public ExecutorService indexTaskThreadPool() {
        ExecutorService threadPool = createThreadPool(INDEX_TASK_THREAD_POOL_NAME);
        AWK_THREAD_POOL.put(INDEX_TASK_THREAD_POOL_NAME, threadPool);
        return threadPool;
    }

    /**
     * 定义一个线程池
     */
    @Bean(name = OTHER_TASK_THREAD_POOL_NAME)
    public ExecutorService otherTaskThreadPool() {
        ExecutorService threadPool = createThreadPool(INDEX_TASK_THREAD_POOL_NAME);
        AWK_THREAD_POOL.put(OTHER_TASK_THREAD_POOL_NAME, threadPool);
        return threadPool;
    }

    @PostConstruct
    public void monitor() {
        MONITOR.scheduleAtFixedRate(() -> {
            for (String threadPoolName : AWK_THREAD_POOL.keySet()) {
                ExecutorService threadPool = AWK_THREAD_POOL.get(threadPoolName);
                if (threadPool instanceof ThreadPoolExecutor threadPoolExecutor) {
                    printThreadPoolInfo(threadPoolName, threadPoolExecutor);
                }
            }
        }, 60, 5, TimeUnit.SECONDS);
    }

    /**
     * 创建线程池
     *
     * @param threadPoolName 线程池名称
     * @return 线程池
     */
    public ExecutorService createThreadPool(String threadPoolName) {
        AwkThreadPoolProperties.PoolConfigItem poolConfigItem = threadPoolProperties.getThreadPoolConfig().get(threadPoolName);
        if (poolConfigItem == null) {
            throw new IllegalArgumentException("ThreadPoolProperties is un config threadPoolName" + threadPoolName);
        }
        RejectedExecutionHandler rejectedExecutionHandler = switch (poolConfigItem.getRejectedPolicy()) {
            case ABORT -> new ThreadPoolExecutor.AbortPolicy();
            case DISCARD -> new ThreadPoolExecutor.DiscardPolicy();
            case CALLER_RUNS, USER_DEFINED -> new ThreadPoolExecutor.CallerRunsPolicy();
            case DISCARD_OLDEST -> new ThreadPoolExecutor.DiscardOldestPolicy();
        };
        NamedThreadFactory namedThreadFactory = new NamedThreadFactory(threadPoolName, poolConfigItem.getDaemon());
        return new ThreadPoolExecutor(poolConfigItem.getCorePoolSize(), poolConfigItem.getMaximumPoolSize(), poolConfigItem.getKeepAliveTime(), TimeUnit.SECONDS, new LinkedBlockingQueue<>(poolConfigItem.getQueueCapacity()), namedThreadFactory, rejectedExecutionHandler);
    }

    private void printThreadPoolInfo(String threadPool, ThreadPoolExecutor threadPoolExecutor) {
        log.info("=====================ThreadPool monitor Info: [{}]=====================", threadPool);
        log.info("Core Pool Size: {}", threadPoolExecutor.getCorePoolSize());
        log.info("Maximum Pool Size: {}", threadPoolExecutor.getMaximumPoolSize());
        log.info("Pool Size: {}", threadPoolExecutor.getPoolSize());
        log.info("Active Count: {}", threadPoolExecutor.getActiveCount());
        log.info("Task Count: {}", threadPoolExecutor.getTaskCount());
        log.info("Completed Task Count: {}", threadPoolExecutor.getCompletedTaskCount());
        log.info("Queue Size: {}", threadPoolExecutor.getQueue().size());
        log.info("Largest Pool Size: {}", threadPoolExecutor.getLargestPoolSize());
        log.info("Keep Alive Time: {}", threadPoolExecutor.getKeepAliveTime(TimeUnit.SECONDS));
        log.info("Queue Remaining Capacity: {}", threadPoolExecutor.getQueue().remainingCapacity());
        log.info("Queue Type: {}", threadPoolExecutor.getQueue().getClass().getSimpleName());
        log.info("Thread Factory Type: {}", threadPoolExecutor.getThreadFactory().getClass().getSimpleName());
        log.info("Rejected Execution Handler Type: {}", threadPoolExecutor.getRejectedExecutionHandler().getClass().getSimpleName());
        log.info("Shutdown Status: {}", threadPoolExecutor.isShutdown());
        log.info("================================================================");
    }
}
